/// Test the load bandwidth of each block in the global memory.
#include <algorithm>
#include <cuda.h>
#include <stdio.h>
#include <stdlib.h>
#include <tuple>
#include "benchmark_config.h"
#include "cudaDeviceQuery.h"
#include "gpu_timer.h"

typedef uint32_t DataType;

constexpr int kIter = 2048;
constexpr int kWarpPerBlock = 16;
constexpr int kMaxStride = 32;

__global__ void mem_bw_8(DataType *array, DataType *sink, int offset = 0,
                         int stride = 1) {
  // thread index
  uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
  DataType temp = 0;
  DataType sum = 0;
  for (int iter = 0; iter < kIter; iter++) {
    temp = array[((iter + 0) * blockDim.x * gridDim.x + tid) * stride + offset];
    sum += temp;
  }
  // Trick compiler not to eliminate LDG, while avoiding uncessary writes.
  if (tid > kWarpSize * kWarpPerBlock) {
    sink[tid] = sum;
  }
}

int main() {

  initializeDeviceProp(0);
  const int BlockNum = 1;
  const int ThreadsPerBlock = kWarpPerBlock * kWarpSize;
  uint64_t total_byte =
      BlockNum * ThreadsPerBlock * sizeof(DataType) * kIter * kMaxStride;

  printf("BlockNum, ThreadsPerBlock: %d, %d\n", BlockNum, ThreadsPerBlock);
  float footprint_MByte = (float)BlockNum * ThreadsPerBlock * sizeof(DataType) *
                          kIter / 1024 / 1024;
  printf("Memory Footprint = %f (MB)\n", footprint_MByte);

  DataType *posArray_g;
  checkGPUErrors(cudaMalloc(&posArray_g, total_byte));
  checkGPUErrors(cudaMemset(posArray_g, 0, total_byte));
  DataType *sink_g;
  checkGPUErrors(cudaMalloc(&sink_g, kWarpSize * sizeof(DataType)));

  // Warpup
  for (int i = 0; i < WARMUPNUM; ++i)
    mem_bw_8<<<BlockNum, ThreadsPerBlock>>>(posArray_g, sink_g);

  int offset = 0;

  GpuTimer timer;
  float total_time = 0.f;

  for (int i = 0; i < 100; ++i) {
    timer.start();
    mem_bw_8<<<BlockNum, ThreadsPerBlock>>>(posArray_g, sink_g, offset);
    timer.stop_and_wait();
    total_time += timer.duration();
  }

  total_time /= 100;

  float bw_sec = footprint_MByte / total_time;
  printf("Average kernel BW: %f GB/s\n", bw_sec);

  checkGPUErrors(cudaFree(posArray_g));
  checkGPUErrors(cudaFree(sink_g));
  return 0;
}
